from pocketflow import Node
from .utils.call_llm import call_llm
from .utils.text_analysis import text_analysis
from .utils.risk_assessment import risk_assessment
from .utils.validate_issues import validate_issues
from .utils.format_correction import format_correction

class FreeConversationNode(Node):
    """
    Handles natural language dialogue in chat mode.
    """
    def prep(self, shared):
        # Read user_input and conversation_history from the shared store
        return {
            "user_input": shared["user_input"],
            "conversation_history": shared.get("conversation_history", [])
        }
    
    def exec(self, data):
        # Prepare conversation history for LLM
        history = data["conversation_history"]
        user_input = data["user_input"]
        
        # Build prompt with conversation history
        prompt = "以下是之前的对话历史（如果有）：\n"
        for i, exchange in enumerate(history):
            prompt += f"用户: {exchange['user']}\n"
            prompt += f"助手: {exchange['assistant']}\n\n"
        
        prompt += f"用户: {user_input}\n"
        prompt += "请以助手的身份回复上述用户的问题。回复应该是有帮助的、礼貌的，并且直接回答用户的问题。"
        
        # Call LLM to generate a response
        response = call_llm(prompt)
        return response
    
    def post(self, shared, prep_res, exec_res):
        # Update conversation history and store response
        if "conversation_history" not in shared:
            shared["conversation_history"] = []
        
        # Add the current exchange to the history
        shared["conversation_history"].append({
            "user": shared["user_input"],
            "assistant": exec_res
        })
        
        # Store the response
        shared["response"] = exec_res
        return "default"

class CorrectionAnalysisNode(Node):
    """
    Analyzes JSON input for issues in check mode.
    """
    def prep(self, shared):
        # Read user_input from the shared store
        return shared["user_input"]
    
    def exec(self, data):
        # Initialize results
        results = {
            "typo_issues": [],
            "semantic_issues": [],
            "risk_level_issues": [],
            "compliance_status_issues": []
        }
        
        # Extract fields from input
        inspector = data.get("inspector", "")
        description = data.get("description", "")
        guide = data.get("guide", "")
        risk_level = data.get("risk_level", "")
        compliance_status = data.get("compliance_status", "")
        title = data.get("title", "")
        
        # Check inspector field for typos and semantic issues
        if inspector:
            text_results = text_analysis(inspector, title)
            results["typo_issues"] = text_results.get("typo_issues", [])
            results["semantic_issues"] = text_results.get("semantic_issues", [])
        
        print("risk_level",risk_level)
        print("compliance_status",compliance_status)
        print("data",data)
        # Evaluate risk level and compliance status
        risk_results = risk_assessment(
            description, 
            guide, 
            inspector, 
            risk_level, 
            compliance_status,
            title
        )
        
        results["risk_level_issues"] = risk_results.get("risk_level_issues", [])
        results["compliance_status_issues"] = risk_results.get("compliance_status_issues", [])
        
        return results
    
    def post(self, shared, prep_res, exec_res):
        # Store analysis results in the shared store
        shared["analysis_results"] = exec_res
        return "default"

class CorrectionReviewNode(Node):
    """
    Reviews and validates correction analysis results.
    """
    def prep(self, shared):
        # Read analysis_results from the shared store
        return shared["analysis_results"]
    
    def exec(self, analysis_results):
        # Validate issues and filter out false positives
        return validate_issues(analysis_results)
    
    def post(self, shared, prep_res, exec_res):
        # Store reviewed_results in the shared store
        shared["reviewed_results"] = exec_res
        return "default"

class ResponseFormattingNode(Node):
    """
    Formats the response according to the mode.
    """
    def prep(self, shared):
        # Read input_type and relevant data from the shared store
        data = {
            "input_type": shared["input_type"]
        }
        
        if shared["input_type"] == "chat":
            data["response"] = shared.get("response", "")
        else:  # check mode
            data["reviewed_results"] = shared.get("reviewed_results", {"valid_issues": [], "is_valid_json": True})
        
        return data
    
    def exec(self, data):
        if data["input_type"] == "chat":
            # For chat mode, use the response directly
            return data["response"]
        else:  # check mode
            # For check mode, format the valid issues
            return format_correction(data["reviewed_results"]["valid_issues"])
    
    def post(self, shared, prep_res, exec_res):
        # Store the final formatted response
        shared["response"] = exec_res
        return "default"